import pandas as pd
from finrl_myself.env.env_portfolio import StockPortfolioEnv
from finrl_myself.env.env_stocktrading import StockTradingEnv
from finrl_myself.ppo.agent import PPO
from finrl_myself.ddpg.agent import DDPG
from finrl_myself.td3.agent import TD3
from finrl_myself.sac.agent import SAC

def train(
        task: str = None,
        env_kwargs: dict = None,
        df_train: pd.DataFrame = None,
        df_validation: pd.DataFrame = None,
        df_trade: pd.DataFrame = None,
        model: str = None,
        agent_kwargs: dict = None,
        training_times: int = None,
):

    if task == 'portfolio_allocation':
        Env = StockPortfolioEnv
    if task == 'stock_trade':
        Env = StockTradingEnv

    models = {'ddpg', 'td3', 'sac', 'ppo'}
    assert model in models, "model is invalid, supported models are: DDPG, TD3, SAC, PPO."
    if model == 'ppo':
        Model = PPO
    elif model == 'ddpg':
        Model = DDPG
    elif model == 'td3':
        Model = TD3
    elif model == 'sac':
        Model = SAC

    for i in range(1, training_times + 1):
        print(f'################################# 这是第{i}次训练 #######################################')
        filename = str(i) + '.pth'
        env_train = Env(df_train, **env_kwargs)
        if df_validation is not None:
            env_validation = Env(df_validation, **env_kwargs)
        else:
            env_validation = None
        env_trade = Env(df_trade, **env_kwargs)
        agent_kwargs.update({'env_train': env_train,
                             'env_validation': env_validation,
                             'env_trade': env_trade,
                             'filename': filename})
        agent = Model(**agent_kwargs)
        agent.train()
        agent.save()


